from abc import ABC

import torch
import numpy as np
from captum.attr import IntegratedGradients
from ts.context import Context
from ts.torch_handler.base_handler import BaseHandler


class Vision3DHandler(BaseHandler, ABC):
    """Refer to ts.torch_handler.vision_handler

    Args:
        BaseHandler (_type_): _description_
        ABC (_type_): _description_
    """
    
    def __init__(self):
        self.model: torch.nn.Module
        super().__init__()

    def initialize(self, context: Context):
        super().initialize(context)
        self.ig = IntegratedGradients(self.model)
        self.initialized = True
        properties = context.system_properties
        if not properties.get("limit_max_image_pixels"):
            # TODO 限制图像大小
            pass
        self.context: Context

    def preprocess(self, data):
        return super().preprocess(data)

    def get_insights(self, tensor_data, _, target=0):
        print("input shape", tensor_data.shape)
        return self.ig.attribute(tensor_data, target=target, n_steps=15).tolist()

